package com.bougaga.opensource.scheduler.netty.scheduler;

import com.bougaga.opensource.scheduler.core.handler.TaskHandler;
import com.bougaga.opensource.scheduler.core.handler.TaskResult;
import com.bougaga.opensource.scheduler.core.scheduler.DelayedTaskScheduler;
import com.bougaga.opensource.scheduler.core.scheduler.SchedulerStatistics;
import com.bougaga.opensource.scheduler.core.task.DelayedTask;
import com.bougaga.opensource.scheduler.core.task.TaskStatus;
import io.netty.util.HashedWheelTimer;
import io.netty.util.Timeout;
import io.netty.util.TimerTask;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import jakarta.annotation.PostConstruct;
import jakarta.annotation.PreDestroy;
import java.time.LocalDateTime;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicLong;
import java.util.stream.Collectors;

/**
 * 基于Netty时间轮的延迟任务调度器
 * 
 * @author bougaga
 * @since 1.0.0
 */
@Component
public class NettyTimeWheelScheduler implements DelayedTaskScheduler {
    
    private static final Logger logger = LoggerFactory.getLogger(NettyTimeWheelScheduler.class);
    
    // 时间轮配置
    private static final long TICK_DURATION = 100; // 100ms per tick
    private static final int WHEEL_SIZE = 512; // 时间轮大小
    
    // 线程池配置
    private static final int CORE_POOL_SIZE = Runtime.getRuntime().availableProcessors();
    private static final int MAX_POOL_SIZE = CORE_POOL_SIZE * 2;
    private static final int QUEUE_CAPACITY = 10000;
    private static final long KEEP_ALIVE_TIME = 60L;
    
    private final HashedWheelTimer timer;
    private final ExecutorService executorService;
    private final AtomicBoolean running = new AtomicBoolean(false);
    private final SchedulerStatistics statistics = new SchedulerStatistics();
    
    // 任务存储和处理器
    private final TaskStorage taskStorage;
    private final Map<String, TaskHandler> taskHandlers;
    
    // 任务执行统计
    private final AtomicLong totalExecutionTime = new AtomicLong(0);
    private final AtomicLong executionCount = new AtomicLong(0);
    
    @Autowired
    public NettyTimeWheelScheduler(TaskStorage taskStorage, List<TaskHandler> handlers) {
        this.taskStorage = taskStorage;
        this.taskHandlers = handlers.stream()
                .collect(Collectors.toMap(TaskHandler::getHandlerName, handler -> handler));
        
        // 初始化时间轮
        this.timer = new HashedWheelTimer(
                Thread::new,
                TICK_DURATION,
                TimeUnit.MILLISECONDS,
                WHEEL_SIZE
        );
        
        // 初始化线程池
        this.executorService = new ThreadPoolExecutor(
                CORE_POOL_SIZE,
                MAX_POOL_SIZE,
                KEEP_ALIVE_TIME,
                TimeUnit.SECONDS,
                new LinkedBlockingQueue<>(QUEUE_CAPACITY),
                new ThreadPoolExecutor.CallerRunsPolicy()
        );
    }
    
    @PostConstruct
    public void init() {
        logger.info("NettyTimeWheelScheduler initialized with tick duration: {}ms, wheel size: {}", 
                TICK_DURATION, WHEEL_SIZE);
    }
    
    @PreDestroy
    public void destroy() {
        stop();
        if (timer != null) {
            timer.stop();
        }
        if (executorService != null) {
            executorService.shutdown();
            try {
                if (!executorService.awaitTermination(60, TimeUnit.SECONDS)) {
                    executorService.shutdownNow();
                }
            } catch (InterruptedException e) {
                executorService.shutdownNow();
                Thread.currentThread().interrupt();
            }
        }
    }
    
    @Override
    public String submitTask(DelayedTask task) {
        if (task == null) {
            throw new IllegalArgumentException("Task cannot be null");
        }
        
        try {
            // 保存任务到存储
            if (!taskStorage.saveTask(task)) {
                throw new RuntimeException("Failed to save task: " + task.getTaskId());
            }
            
            // 计算延迟时间
            long delay = task.getDelayTimeMillis();
            if (delay <= 0) {
                // 立即执行
                executeTask(task);
            } else {
                // 延迟执行
                scheduleTask(task, delay);
            }
            
            // 更新统计信息
            statistics.setTotalTasks(statistics.getTotalTasks() + 1);
            statistics.setPendingTasks(statistics.getPendingTasks() + 1);
            
            logger.info("Task submitted successfully: {}, delay: {}ms", task.getTaskId(), delay);
            return task.getTaskId();
        } catch (Exception e) {
            logger.error("Failed to submit task: {}", task.getTaskId(), e);
            throw new RuntimeException("Failed to submit task", e);
        }
    }
    
    @Override
    public String submitTask(String taskName, String taskType, String taskData, Long delayMillis) {
        DelayedTask task = new DelayedTask(taskName, taskType, taskData, delayMillis);
        return submitTask(task);
    }
    
    @Override
    public String submitTask(String taskName, String taskType, String taskData, LocalDateTime executeTime) {
        DelayedTask task = new DelayedTask(taskName, taskType, taskData, executeTime);
        return submitTask(task);
    }
    
    @Override
    public boolean cancelTask(String taskId) {
        try {
            DelayedTask task = taskStorage.getTask(taskId);
            if (task == null) {
                return false;
            }
            
            if (task.getStatus().isCompleted()) {
                return false;
            }
            
            // 更新任务状态
            task.setStatus(TaskStatus.CANCELLED);
            taskStorage.updateTask(task);
            
            // 更新统计信息
            statistics.setCancelledTasks(statistics.getCancelledTasks() + 1);
            if (task.getStatus() == TaskStatus.PENDING) {
                statistics.setPendingTasks(statistics.getPendingTasks() - 1);
            } else if (task.getStatus() == TaskStatus.RUNNING) {
                statistics.setRunningTasks(statistics.getRunningTasks() - 1);
            }
            
            logger.info("Task cancelled successfully: {}", taskId);
            return true;
        } catch (Exception e) {
            logger.error("Failed to cancel task: {}", taskId, e);
            return false;
        }
    }
    
    @Override
    public DelayedTask getTask(String taskId) {
        return taskStorage.getTask(taskId);
    }
    
    @Override
    public List<DelayedTask> getAllTasks() {
        return taskStorage.getAllTasks();
    }
    
    @Override
    public List<DelayedTask> getPendingTasks() {
        return taskStorage.getPendingTasks();
    }
    
    @Override
    public List<DelayedTask> getRunningTasks() {
        return taskStorage.getTasksByStatus(TaskStatus.RUNNING);
    }
    
    @Override
    public List<DelayedTask> getCompletedTasks() {
        List<DelayedTask> completedTasks = new ArrayList<>();
        completedTasks.addAll(taskStorage.getTasksByStatus(TaskStatus.SUCCESS));
        completedTasks.addAll(taskStorage.getTasksByStatus(TaskStatus.FAILED));
        completedTasks.addAll(taskStorage.getTasksByStatus(TaskStatus.CANCELLED));
        completedTasks.addAll(taskStorage.getTasksByStatus(TaskStatus.TIMEOUT));
        return completedTasks;
    }
    
    @Override
    public List<DelayedTask> getTasksByType(String taskType) {
        return taskStorage.getTasksByType(taskType);
    }
    
    @Override
    public List<DelayedTask> getTasksByStatus(String status) {
        try {
            TaskStatus taskStatus = TaskStatus.valueOf(status.toUpperCase());
            return taskStorage.getTasksByStatus(taskStatus);
        } catch (IllegalArgumentException e) {
            logger.warn("Invalid task status: {}", status);
            return Collections.emptyList();
        }
    }
    
    @Override
    public void start() {
        if (running.compareAndSet(false, true)) {
            logger.info("NettyTimeWheelScheduler started");
            
            // 启动定时任务，定期检查过期任务
            timer.newTimeout(new TimerTask() {
                @Override
                public void run(Timeout timeout) throws Exception {
                    if (running.get()) {
                        checkExpiredTasks();
                        // 重新调度
                        timer.newTimeout(this, 1, TimeUnit.SECONDS);
                    }
                }
            }, 1, TimeUnit.SECONDS);
        }
    }
    
    @Override
    public void stop() {
        if (running.compareAndSet(true, false)) {
            logger.info("NettyTimeWheelScheduler stopped");
        }
    }
    
    @Override
    public boolean isRunning() {
        return running.get();
    }
    
    @Override
    public SchedulerStatistics getStatistics() {
        // 更新统计信息
        statistics.setPendingTasks(taskStorage.getTaskCountByStatus(TaskStatus.PENDING));
        statistics.setRunningTasks(taskStorage.getTaskCountByStatus(TaskStatus.RUNNING));
        statistics.setSuccessTasks(taskStorage.getTaskCountByStatus(TaskStatus.SUCCESS));
        statistics.setFailedTasks(taskStorage.getTaskCountByStatus(TaskStatus.FAILED));
        statistics.setCancelledTasks(taskStorage.getTaskCountByStatus(TaskStatus.CANCELLED));
        statistics.setRetryTasks(taskStorage.getTaskCountByStatus(TaskStatus.RETRYING));
        statistics.setTimeoutTasks(taskStorage.getTaskCountByStatus(TaskStatus.TIMEOUT));
        statistics.setTotalTasks(taskStorage.getAllTasks().size());
        // 计算平均执行时间
        long count = executionCount.get();
        if (count > 0) {
            statistics.setAverageExecutionTime((double) totalExecutionTime.get() / count);
        }
        
        statistics.update();
        return statistics;
    }
    
    @Override
    public int cleanupCompletedTasks(LocalDateTime beforeTime) {
        int cleanedCount = taskStorage.cleanupCompletedTasks(beforeTime);
        logger.info("Cleaned up {} completed tasks", cleanedCount);
        return cleanedCount;
    }
    
    /**
     * 调度任务
     */
    private void scheduleTask(DelayedTask task, long delay) {
        timer.newTimeout(new TimerTask() {
            @Override
            public void run(Timeout timeout) throws Exception {
                if (running.get()) {
                    executeTask(task);
                }
            }
        }, delay, TimeUnit.MILLISECONDS);
    }
    
    /**
     * 执行任务
     */
    private void executeTask(DelayedTask task) {
        if (!running.get()) {
            return;
        }
        
        // 检查任务状态
        DelayedTask currentTask = taskStorage.getTask(task.getTaskId());
        if (currentTask == null || !currentTask.getStatus().canExecute()) {
            return;
        }
        
        // 更新任务状态为执行中
        currentTask.setStatus(TaskStatus.RUNNING);
        taskStorage.updateTask(currentTask);
        
        // 更新统计信息
        statistics.setPendingTasks(statistics.getPendingTasks() - 1);
        statistics.setRunningTasks(statistics.getRunningTasks() + 1);
        
        // 异步执行任务
        executorService.submit(() -> {
            long startTime = System.currentTimeMillis();
            TaskResult result = null;
            
            try {
                // 获取任务处理器
                TaskHandler handler = getTaskHandler(currentTask);
                if (handler == null) {
                    result = TaskResult.failure("No handler found for task type: " + currentTask.getTaskType());
                } else {
                    // 执行任务
                    handler.beforeExecute(currentTask);
                    result = handler.execute(currentTask);
                    handler.afterExecute(currentTask, result);
                }
            } catch (Exception e) {
                logger.error("Task execution failed: {}", currentTask.getTaskId(), e);
                result = TaskResult.failure("Task execution failed", e);
                
                // 调用异常回调
                TaskHandler handler = getTaskHandler(currentTask);
                if (handler != null) {
                    handler.onException(currentTask, e);
                }
            } finally {
                // 更新任务状态
                updateTaskStatus(currentTask, result);
                
                // 更新统计信息
                long executionTime = System.currentTimeMillis() - startTime;
                totalExecutionTime.addAndGet(executionTime);
                executionCount.incrementAndGet();
                
                if (result != null && result.isSuccess()) {
                    statistics.setSuccessTasks(statistics.getSuccessTasks() + 1);
                } else {
                    statistics.setFailedTasks(statistics.getFailedTasks() + 1);
                }
                
                statistics.setRunningTasks(statistics.getRunningTasks() - 1);
                
                // 记录执行时间统计
                if (executionTime > statistics.getMaxExecutionTime()) {
                    statistics.setMaxExecutionTime(executionTime);
                }
                if (statistics.getMinExecutionTime() == 0 || executionTime < statistics.getMinExecutionTime()) {
                    statistics.setMinExecutionTime(executionTime);
                }
            }
        });
    }
    
    /**
     * 获取任务处理器
     */
    private TaskHandler getTaskHandler(DelayedTask task) {
        if (task.getHandlerName() != null) {
            return taskHandlers.get(task.getHandlerName());
        }
        
        // 根据任务类型查找处理器
        return taskHandlers.values().stream()
                .filter(handler -> handler.supports(task.getTaskType()))
                .findFirst()
                .orElse(null);
    }
    
    /**
     * 更新任务状态
     */
    private void updateTaskStatus(DelayedTask task, TaskResult result) {
        try {
            if (result != null && result.isSuccess()) {
                task.setStatus(TaskStatus.SUCCESS);
            } else {
                if (task.canRetry()) {
                    task.setStatus(TaskStatus.RETRYING);
                    task.incrementRetryCount();
                    
                    // 计算重试时间
                    long retryDelay = task.getRetryInterval();
                    if (retryDelay > 0) {
                        task.setExecuteTime(LocalDateTime.now().plusNanos(retryDelay * 1_000_000));
                        scheduleTask(task, retryDelay);
                    }
                } else {
                    task.setStatus(TaskStatus.FAILED);
                }
            }
            
            taskStorage.updateTask(task);
        } catch (Exception e) {
            logger.error("Failed to update task status: {}", task.getTaskId(), e);
        }
    }
    
    /**
     * 检查过期任务
     */
    private void checkExpiredTasks() {
        try {
            List<DelayedTask> expiredTasks = taskStorage.getExpiredTasks();
            for (DelayedTask task : expiredTasks) {
                if (task.getStatus() == TaskStatus.PENDING) {
                    executeTask(task);
                }
            }
        } catch (Exception e) {
            logger.error("Failed to check expired tasks", e);
        }
    }
    
    /**
     * 任务存储接口
     */
    public interface TaskStorage {
        boolean saveTask(DelayedTask task);
        DelayedTask getTask(String taskId);
        boolean updateTask(DelayedTask task);
        boolean deleteTask(String taskId);
        List<DelayedTask> getAllTasks();
        List<DelayedTask> getTasksByStatus(TaskStatus status);
        List<DelayedTask> getTasksByType(String taskType);
        List<DelayedTask> getPendingTasks();
        List<DelayedTask> getExpiredTasks();
        long getTaskCountByStatus(TaskStatus status);
        int cleanupCompletedTasks(LocalDateTime beforeTime);
    }
}
